import asyncio
import base64
import hashlib
import hmac
import json
from urllib.parse import urlparse, urlencode
from datetime import datetime
from time import mktime
from wsgiref.handlers import format_date_time
import websockets


class SparkAPI:
    def __init__(self, app_id, api_key, api_secret, spark_url, domain):
        self.app_id = app_id
        self.api_key = api_key
        self.api_secret = api_secret
        self.spark_url = spark_url
        self.domain = domain

    class WsParam:
        def __init__(self, APPID, APIKey, APISecret, Spark_url):
            self.APPID = APPID
            self.APIKey = APIKey
            self.APISecret = APISecret
            self.host = urlparse(Spark_url).netloc
            self.path = urlparse(Spark_url).path
            self.Spark_url = Spark_url

        def create_url(self):
            now = datetime.now()
            date = format_date_time(mktime(now.timetuple()))
            signature_origin = f"host: {self.host}\n"
            signature_origin += f"date: {date}\n"
            signature_origin += f"GET {self.path} HTTP/1.1"
            signature_sha = hmac.new(
                self.APISecret.encode("utf-8"),
                signature_origin.encode("utf-8"),
                digestmod=hashlib.sha256,
            ).digest()
            signature_sha_base64 = base64.b64encode(signature_sha).decode("utf-8")
            authorization_origin = (
                f'api_key="{self.APIKey}", algorithm="hmac-sha256", '
                f'headers="host date request-line", signature="{signature_sha_base64}"'
            )
            authorization = base64.b64encode(authorization_origin.encode("utf-8")).decode(
                "utf-8"
            )
            v = {"authorization": authorization, "date": date, "host": self.host}
            url = self.Spark_url + "?" + urlencode(v)
            return url

    def gen_params(self, appid, question, system_prompt):
        """生成请求参数，包含系统角色设定"""
        data = {
            "header": {"app_id": appid, "uid": "1234"},
            "parameter": {
                "chat": {"domain": self.domain, "temperature": 0.5, "max_tokens": 2048}
            },
            "payload": {
                "message": {
                    "text": [
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": question},
                    ]
                }
            },
        }
        return data

    async def ask_question(self, question, system_prompt=""):
        """向讯飞星火API发送问题并返回回答"""
        ws_param = self.WsParam(self.app_id, self.api_key, self.api_secret, self.spark_url)
        ws_url = ws_param.create_url()
        answer = ""

        async with websockets.connect(ws_url) as ws:
            # 发送问题
            data = json.dumps(self.gen_params(self.app_id, question, system_prompt))
            await ws.send(data)

            while True:
                try:
                    message = await ws.recv()
                    data = json.loads(message)
                    code = data["header"]["code"]
                    if code != 0:
                        print(f"请求出错: {code}, {data}")
                        break
                    else:
                        choices = data["payload"]["choices"]
                        status = choices["status"]
                        content = choices["text"][0]["content"]
                        answer += content
                        if status == 2:  # 会话结束
                            break
                except websockets.ConnectionClosed:
                    print("连接关闭")
                    break

        return answer


# === 外部调用函数 ===
def ai_reply(question):
    APP_ID = "dc80c7ca"
    API_KEY = "f2f3b2100524a338af238197ac03274e"
    API_SECRET = "MGIzYmQ1YTRlYTI0NjZhMzc2MGI1MGRm"

    DOMAIN = "xop3qwen1b7"
    SPARK_URL = "wss://maas-api.cn-huabei-1.xf-yun.com/v1.1/chat"

    spark_api = SparkAPI(APP_ID, API_KEY, API_SECRET, SPARK_URL, DOMAIN)

    custom_prompt = """你是一个安装在用户家门口的"智能门神"助手，具有以下特点和功能：

    角色定位：
    1. 你是用户的家庭智能助手，负责迎宾、提醒和提供信息
    2. 你的语气应该友好、亲切但专业，像一个贴心的管家
    3. 你主要使用中文进行交流，回复简洁明了（1-3句话）

    简单对话：进行友好的日常交流，讲笑话、提供建议等
    行为准则：
    . 回复简洁，每次回答尽量控制在1-3句话内
    . 当不确定如何回答时，引导用户询问其他你能回答的问题
    请根据以上设定，以智能门神的身份回应用户的请求。"""
    # 创建新的事件循环
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        result = loop.run_until_complete(
            spark_api.ask_question(question, custom_prompt)
        )
        return result
    finally:
        loop.close()
